import pandas as pd
import numpy as np
import math
import glob
from tqdm import tqdm





height = 335
buildings_dir = "H:/data/zzdxOSDB/buildings/20240116-LOD/"
model_center_txt_file = buildings_dir + "models_center.txt"
out_pos_file = buildings_dir + str(height) + "/train_cam_pos.h335.y45.p60.2.txt"

scene_id = 1

def get_instance_centerpoints(allpoints):
    sem_ins_labels = np.ascontiguousarray(allpoints[:, -2:])
    sem_ins_labels = sem_ins_labels.astype(np.int32)
    
    building_instance_labels = np.where(sem_ins_labels[:,0] == 1,sem_ins_labels[:,1],-100)
    building_instance_labels = np.unique(building_instance_labels)[1:]

    # building_instance_points_index = np.where(sem_ins_labels[:,0] == 1)

    ins_centre_xyz = {}
    for id in building_instance_labels:
        index = np.where(sem_ins_labels[:,1] == id)
        points = np.ascontiguousarray(allpoints[index])
        xyz = np.mean(points[:,:3],axis=0)
        ins_centre_xyz[id] = xyz


    return ins_centre_xyz

def get_cam_pos(center,height,heading,pitch):
    
    heading = heading/180 * math.pi
    pitch = pitch/180 * math.pi
    #abs_h = height - center[2]
    abs_h = height

    xy_r = math.tan(math.pi/2 + pitch) * abs_h * -1
    cam_y = xy_r * math.cos(heading)
    cam_x = xy_r * math.sin(heading)

    cam_x += center[0]
    cam_y += center[1]

    return cam_x,cam_y


def gen_pos(ins_centre_xyz):
    # heading = [0,60,120,180,240,300]
    # heading = [0,45,90,135,180,225,270,315]
    heading = [22.5,67.5,112.5,157.5,202.5,247.5,292.5,337.5]


    pitch = [-60]

    res_cam_pose = []
    for r in range(ins_centre_xyz.shape[0]):
    # for id,xyz in ins_centre_xyz.items():
        rd = ins_centre_xyz[r]
        id = int(rd[0])
        xyz = rd[1:]
        for p in pitch:
            for h in heading:
                cx,cy = get_cam_pos(xyz,height,h,p)
                res_cam_pose.append([cx,cy,height,h,p,0,scene_id,id])
    return res_cam_pose

def write_file(out_pos_file_path,res_cam_pose):

    out_pos_file = open(out_pos_file_path,'w+')
    out_pos_file.write("#x y z heading pitch roll scene_id id\n")
    for pose in res_cam_pose: 
        str2 = '@'.join([str(f) for f in pose])
        out_pos_file.write(str2+"\n")


def gen_pose_file(path,outpath):
    points = pd.read_csv(path, header=None).values
    cam_pose = gen_pos(points)
    write_file(outpath,cam_pose)



gen_pose_file(model_center_txt_file,out_pos_file)



